"""
进入该文件所在的目录下，执行 python f0.py即可执行。（代码使用python3)
算法的参数 设置在代码的第232-242行，f0()函数的前10行。
"""
import matplotlib.pyplot as plt
import numpy as np
import scipy.io.wavfile as siowav
from numpy.lib import stride_tricks
from scipy.signal import butter, lfilter


# ===== low pass filter  以下两个函数为一个带通滤滤波器 ==========

def butter_bandpass(lowcut, highcut, fs, order=5):
    nyq = 0.5 * fs
    low = lowcut / nyq
    high = highcut / nyq
    b, a = butter(order, [low, high], btype='band')
    return b, a


def butter_bandpass_filter(data, lowcut, highcut, fs, order=5):
    b, a = butter_bandpass(lowcut, highcut, fs, order=order)
    y = lfilter(b, a, data)
    return y


# ===== Load data ==============

def load_wav(wav_file, is_norm=True):
    """加载数据"""
    sr, signals = siowav.read(wav_file)
    signals = signals.astype(float)
    if is_norm:
        signals = signals / np.max(np.abs(signals))
    return sr, signals


def get_frames(signals, frame_size, hop_size):
    """
    把语音信号转换为易于numpy处理的2dArray.
    输出一个矩阵，每一行代表每一帧要处理的数据，行数代表帧数。
    """
    frame_num = np.floor((len(signals) - frame_size) / float(hop_size)).astype(int) + 1
    ele_stride = signals.strides[-1]
    frames = stride_tricks.as_strided(signals, shape=(frame_num, frame_size),
                                      strides=(ele_stride * hop_size, ele_stride)).copy()
    return frames


# ============= Window function 四种窗函数的实现 =================================

def rect_window(n):
    return np.ones(n)


def gaussian_window(n):
    sigma = 0.5
    power = -0.5 * ((np.arange(n) - (n - 1) / 2) / (sigma * (n - 1) / 2)) ** 2
    return np.exp(power)


def hamming_window(n):
    return 0.54 - 0.46 * np.cos(2 * np.pi * np.arange(n) / (n - 1))


def hann_window(n):
    return 0.5 * (1 - np.cos(2 * np.pi * np.arange(n) / (n - 1)))


# ======= Calculate F0 ==============

def calcR(frames, window_func):
    """计算短时自相关函数"""
    frame_num, frame_size = frames.shape
    wined_frames = frames * window_func(frame_size)
    R = np.zeros((frame_num, frame_size))
    for i in range(1, frame_size):
        R[:, i] = np.sum(wined_frames[:, :-i] * wined_frames[:, i:], axis=1)
    R[:, 0] = np.sum(wined_frames ** 2, axis=1)
    return R


def get_avail_parts(R, R0_min):
    """根据短时自相关函数的值和一个阈值，求出存在基音的部分"""
    points = list()
    point = np.array([0, 0], dtype=np.int32)
    status = 0
    R0 = R[:, 0]
    for i, ele in zip(range(len(R0)), R0):
        if ele >= R0_min and status == 0:
            status = 1
            point[0] = i
        elif ele < R0_min and status == 1:
            status = 0
            point[1] = i
            points.append(point.copy())
    return points


def getF0_coeff(R, R0_min):
    """
    根据短时自相关函数的值和一个阈值，求出存在基音的部分
    返回一个长度等于帧数的向量，若某帧不存在基音，则对应的元素为0， 反之，对应的元素为1。
    """
    return np.where(R[:, 0] >= R0_min, 1, 0)


def get_cycle(R, t1, t2):
    """
    根据短时自相关函数求基音周期
    """
    cycle = np.argmax(R[:, t1:t2], axis=1) + t1
    return cycle


def getF0s(cycles, sr):
    """
    根据基音周期求基频
    """
    cycles = cycles.astype(float)
    F0s = sr / cycles
    return F0s


def smooth_F0_diff(points, F0s, diff, len_limit):
    """
    去除基频曲线中的异常点
    """
    new_points = list()
    new_point = np.array([0, 0])
    for point in points:
        start_i = int(point[0])
        end_i = int(point[1])
        while start_i < end_i and abs(F0s[start_i] - F0s[start_i + 1]) > diff:
            start_i += 1
        new_point[0] = start_i
        for k in range(start_i + 1, end_i):
            if abs(F0s[k] - F0s[k - 1]) > diff:
                new_point[1] = k
                if new_point[1] - new_point[0] > len_limit:
                    new_points.append(new_point.copy())
                new_point[0] = k + 1
        new_point[1] = end_i
        if new_point[1] - new_point[0] > len_limit:
            new_points.append(new_point.copy())
    return new_points


# ========= pulse ============
def get_cycle_from_sample_point(cur_pos, avail_cycles, frame_size, hop_size):
    """
    输入任一个采样点的位置（在signals中的下标），求该采样点对应的基音周期。（画pulse竖线时用到该函数）
    """
    if cur_pos < frame_size / 2:
        return 0
    else:
        idx = int((cur_pos - frame_size / 2) / hop_size)
        if idx > len(avail_cycles) - 1:
            return 0
        return avail_cycles[idx]


# ========= draw ===========

def draw(sr, signals, t, F0s, points, avail_cycles, min_f0, max_f0, frame_size, hop_size,
         x_lim=[0, 13]):
    """画出Pulse竖线和基频曲线"""
    plt.subplot(2, 1, 1)
    plt.xlim(x_lim)
    plt.ylabel('pulse')
    cur_pos = 0
    while cur_pos < len(signals):
        cur_cycle = get_cycle_from_sample_point(cur_pos, avail_cycles, frame_size, hop_size)
        if cur_cycle == 0:
            cur_pos += 1
        else:
            plt.axvline(cur_pos / sr, color='b')
            cur_pos += cur_cycle
    plt.plot(1 / sr * np.arange(len(signals)), signals, color='r')
    plt.subplot(2, 1, 2)
    plt.ylim([min_f0, max_f0])
    plt.xlim(x_lim)
    plt.ylabel('F0')
    plt.xlabel('t')
    for point in points:
        start_i = int(point[0])
        end_i = int(point[1])
        plt.plot(t[start_i: end_i], F0s[start_i: end_i], color='b')
    plt.show()


def draw_pulse(sr, signals, t, F0s, points, avail_cycles, min_f0, max_f0, frame_size, hop_size,
               x_lim=None):
    """只画pulse竖线"""
    plt.plot(1 / sr * np.arange(len(signals)), signals, color='r')
    if x_lim is None:
        plt.xlim([0, 13])
    else:
        plt.xlim(x_lim)
    plt.ylabel('pulse')
    cur_pos = 0
    while cur_pos < len(signals):
        cur_cycle = get_cycle_from_sample_point(cur_pos, avail_cycles, frame_size, hop_size)
        if cur_cycle == 0:
            cur_pos += 1
        else:
            plt.axvline(cur_pos / sr, color='b')
            cur_pos += cur_cycle
    plt.show()


def draw_f0(sr, signals, t, F0s, points, avail_cycles, min_f0, max_f0, frame_size, hop_size,
            y_lim=None):
    """只画基频曲线"""
    if y_lim is None:
        plt.ylim([min_f0, max_f0])
    else:
        plt.ylim(y_lim)
    plt.xlim([0, 13])
    plt.ylabel('F0')
    plt.xlabel('t')
    for point in points:
        start_i = int(point[0])
        end_i = int(point[1])
        plt.plot(t[start_i: end_i], F0s[start_i: end_i], color='b')
    plt.show()


# ======== main ===========
def f0():
    wav_file = './connected.wav'
    window_len = 0.02
    hop_len = 0.005
    min_f0 = 70
    max_f0 = 300
    R0_min = 0.05
    low_cut = 60
    high_cut = 900
    F0_diff = 30
    # draw_type = 0, 既画pulse又画F0; draw_type = 1, 只画pulse; draw_type = 2, 只画F0
    draw_type = 1

    point_len_limit = 3
    sr, signals = load_wav(wav_file)
    frame_size = int(sr * window_len)
    hop_size = int(sr * hop_len)
    t1 = int(sr / max_f0)
    t2 = int(sr / min_f0)
    signals = butter_bandpass_filter(signals, low_cut, high_cut, sr)
    frames = get_frames(signals, frame_size, hop_size)
    R = calcR(frames, window_func=hamming_window)
    # R_sum = np.sum(R, axis=1)
    cycles = get_cycle(R, t1, t2)
    F0s = getF0s(cycles, sr)
    avail_cycles = cycles * getF0_coeff(R, R0_min)
    avail_points = get_avail_parts(R, R0_min)
    smooth_avail_points = smooth_F0_diff(avail_points, F0s, F0_diff, point_len_limit)
    F0_t = window_len / 2 + np.arange(frames.shape[0]) * hop_len
    if draw_type == 0:
        draw(sr, signals, F0_t, F0s, smooth_avail_points, avail_cycles, min_f0, max_f0, frame_size,
             hop_size)
    elif draw_type == 1:
        draw_pulse(sr, signals, F0_t, F0s, smooth_avail_points, avail_cycles, min_f0, max_f0,
                   frame_size,
                   hop_size,
                   x_lim=[7.9, 8.0])
    elif draw_type == 2:
        draw_f0(sr, signals, F0_t, F0s, smooth_avail_points, avail_cycles, min_f0, max_f0,
                frame_size,
                hop_size,
                y_lim=[100, 200])


if __name__ == '__main__':
    f0()
